import Parser

clean_methods = {}
def clean_AST(self):
    return clean_methods[self.type](self)

Parser.Parsed.clean = clean_AST

def clean(aName):    
    def aDecorator(aMethod):
        clean_methods[aName] = aMethod
    return aDecorator

@clean('Statement_0')
def simple(self):
    return AST('FunctionDecl', 
        Parameters = self.nodes[2].clean(), 
        Code = self.nodes[4].clean(),
    )

@clean('Statement_1')
def simple(self):
    return self.nodes[0]

@clean('Statement_2')
def simple(self):
    return AST('Assignment', 
        Var = self.nodes[0], 
        Expression = self.nodes[2].clean(), 
    )

@clean('Statement_3')
def simple(self):
    return AST('Call', 
        Function = self.nodes[0],
        Arguments = [ expression.clean() for expression in 
self.nodes[1].nodes ],
    )
	
@clean('ParList_0')
def simple(self):
    par_list = [ self.nodes[0] ]
    for par_item in self.nodes[1].nodes:
        par_list += [ par_item.nodes[1] ]
    return par_list

@clean('ParList_1')
def simple(self):
    return []

@clean('ArgList_0')
def simple(self):
    arg_list = [ self.nodes[0] ]
    for arg_item in self.nodes[1].nodes:
        arg_list += [ arg_item.nodes[1] ]
    return arg_list

@clean('ArgList_1')
def simple(self):
    return []

@clean('NewScope')
def simple(self):
    return AST('NewScope', Statements=self.nodes[1].clean() )

@clean('Statement__many_0')
def simple(self):
    return [ node.clean() for node in self.nodes ]

@clean('Expression_0')
def simple(self):
    return AST('Call',
        Function = self.nodes[0],
        Arguments = self.nodes[2].clean(),
    )

def first_node(self):
    return self.nodes[0].clean()

for type in ['Statement', 'ParList', 'Expression', 'ArgList']:
    clean_methods[type] = first_node

def no_clean(self):
    return self

for type in ['Var', 'Value']:
    clean_methods[type] = no_clean
